package main

import (
	"io"
	"log"
	"net"
	"net/http"

	"github.com/gorilla/websocket"
)

// WebSocket 升级器（允许跨域）
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 后端 TCP 服务器地址（你的聊天服务器）
const tcpServerAddr = "localhost:8888"

// 处理 WebSocket 连接，转发数据到 TCP 服务器
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 升级 HTTP 连接为 WebSocket
	wsConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket 升级失败:", err)
		return
	}
	defer wsConn.Close()

	// 连接后端 TCP 服务器
	tcpConn, err := net.Dial("tcp", tcpServerAddr)
	if err != nil {
		log.Println("连接 TCP 服务器失败:", err)
		wsConn.WriteMessage(websocket.TextMessage, []byte(`{"type":4,"content":"无法连接到聊天服务器"}`))
		return
	}
	defer tcpConn.Close()

	// 从 WebSocket 读数据，转发到 TCP
	go func() {
		for {
			_, msg, err := wsConn.ReadMessage()
			if err != nil {
				log.Println("WebSocket 读取失败:", err)
				return
			}
			// 转发到 TCP 服务器
			_, err = tcpConn.Write(msg)
			if err != nil {
				log.Println("TCP 写入失败:", err)
				return
			}
		}
	}()

	// 从 TCP 读数据，转发到 WebSocket
	buf := make([]byte, 4096)
	for {
		n, err := tcpConn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Println("TCP 读取失败:", err)
			}
			return
		}
		// 转发到 WebSocket 前端
		err = wsConn.WriteMessage(websocket.TextMessage, buf[:n])
		if err != nil {
			log.Println("WebSocket 写入失败:", err)
			return
		}
	}
}

// 在 bridge.go 的 main 函数中修改
func main() {
	// 托管前端页面（假设前端文件在 bridge 文件夹的 public 子目录）
	http.Handle("/", http.FileServer(http.Dir("./public")))
	// 保留 WebSocket 桥接逻辑
	http.HandleFunc("/tcp-bridge", handleWebSocket)

	log.Println("桥接服务启动，监听 http://localhost:8081")
	log.Fatal(http.ListenAndServe(":8081", nil))
}
